from flask import Flask, request, jsonify, send_from_directory
from langchain.chains import create_sql_query_chain
from langchain.prompts import PromptTemplate
from utiles import extract_sql_from_markdown, load_table_info_from_json
from oracle_connection import get_db_connection, FewShotPromptTemplate
from generate_echarts import visualize_sql_result, analyze_data_for_visualization, convert_sql_result_to_dataframe
from llm_module import SiliconFlow
import os
import traceback

import datetime
import decimal
import json


def serialize_row(row):
    """将SQLAlchemy Row对象转换为可序列化的字典"""
    if row is None:
        return None

        # 如果已经是字典类型，则直接返回
    if isinstance(row, dict):
        return {k: serialize_row(v) for k, v in row.items()}

        # 如果是列表，递归处理每一项
    if isinstance(row, (list, tuple)):
        return [serialize_row(item) for item in row]

        # 处理Row对象
    if hasattr(row, '_mapping'):  # SQLAlchemy 2.0
        return {k: serialize_row(v) for k, v in dict(row._mapping).items()}
    elif hasattr(row, '_asdict') and callable(getattr(row, '_asdict')):
        return {k: serialize_row(v) for k, v in row._asdict().items()}
    elif hasattr(row, '_fields'):  # namedtuple类型
        return {k: serialize_row(getattr(row, k)) for k in row._fields}
    elif hasattr(row, '__dict__') and not isinstance(row, type):
        return {k: serialize_row(v) for k, v in row.__dict__.items() if not k.startswith('_')}

        # 处理特殊类型
    if isinstance(row, (datetime.date, datetime.datetime)):
        return row.isoformat()
    elif isinstance(row, decimal.Decimal):
        return float(row)
    elif isinstance(row, (bytes, bytearray)):
        return row.hex()
    elif isinstance(row, (int, float, str, bool, type(None))):
        return row
    else:
        return str(row)  # 最后手段：转换为字符串


# 自定义JSONEncoder类（用于调试）
class CustomJSONEncoder(json.JSONEncoder):
    def default(self, obj):
        return serialize_row(obj)

# 创建Flask应用并注册编码器
app = Flask(__name__)
app.json_encoder = CustomJSONEncoder

# 初始化全局变量
llm = SiliconFlow()

# 加载全部表信息
full_table_info = load_table_info_from_json("oracle_table_info.json")
print(f"全部表信息中包含 {len(full_table_info)} 个表")

# 指定要包含的表
include_tables = [
    'BAYONET_BUSSINESS.BOAT_WARNING',
    'BAYONET_BUSSINESS.OFF_SITE_CASE',
    'BAYONET_BUSSINESS.ZD_WARNING_TYPE',
    'BAYONET_DYNAMIC.VIDEO_ANALYSIS',
    'BAYONET_DYNAMIC.DATAFUSION',
    'BAYONET_BASICS.SYS_BAYONET',
    'BAYONET_BUSSINESS.MESSAGE_SEND_RECORD'
]

# 获取数据库连接
db = get_db_connection(
    sample_rows=2,
    custom_table_info=full_table_info,
    include_tables=include_tables
)

# 配置SQL生成提示模板
SQL_PROMPT = PromptTemplate.from_template(
    """你是一个SQL专家。根据以下表结构：  
    {table_info}  

    问题：{input}  
    请严格按以下规则响应：  
    1. 只输出SQL代码  
    2. 使用标准Oracle SQL语法  
    3. 不要包含分号结尾  
    4. 最多返回{top_k}条记录  
    5. Oracle中使用ROWNUM来限制返回记录数  

    SQL查询："""
)

# 示例配置
examples = [
    {
        "input": "查询名称中包含'船闸'的所有卡口基本信息，按名称排序。",
        "query": "SELECT ID, NAME, UNIT_ID, DES, PHOTOHTTP FROM BAYONET_BASICS.SYS_BAYONET WHERE NAME LIKE '%船闸%' ORDER BY NAME",
    },
    {
        "input": "统计一月份各卡口点的船舶通过数量，按照通过船舶数量降序排列，只显示通过量前10的卡口。",
        "query": "SELECT * FROM (SELECT b.NAME AS 卡口名称, COUNT(d.CODE) AS 通过船舶数量 FROM BAYONET_BASICS.SYS_BAYONET b JOIN BAYONET_DYNAMIC.DATAFUSION d ON b.ID = d.BAYONET_ID WHERE d.PASSTIME >= TRUNC(SYSDATE, 'YYYY') AND d.PASSTIME < ADD_MONTHS(TRUNC(SYSDATE, 'YYYY'), 1) GROUP BY b.NAME ORDER BY 通过船舶数量 DESC) WHERE ROWNUM <= 10",
    },
    {
        "input": "查询系统中各类型预警的发生次数和占比，按发生次数从高到低排序。",
        "query": "WITH warning_counts AS (SELECT t.NAME AS 预警类型, COUNT(w.CODE) AS 预警次数 FROM BAYONET_BUSSINESS.ZD_WARNING_TYPE t JOIN BAYONET_BUSSINESS.BOAT_WARNING w ON t.ID = w.WARNING_TYPE GROUP BY t.NAME) SELECT 预警类型,预警次数, ROUND(预警次数 * 100 / SUM(预警次数) OVER(), 2) || '%' AS 占比 FROM warning_counts ORDER BY 预警次数 DESC"
    },
    {
        "input": "查询名为'苏无锡货08168'的船舶最近一次通过系统的完整记录，包括通过卡口信息、预警情况和消息发送记录。",
        "query": "SELECT * FROM (SELECT d.PASSTIME AS 通过时间, b.NAME AS 卡口名称, d.BOAT_NAME AS 船舶名称, wt.NAME AS 预警类型, w.WARNING_TIME AS 预警时间, w.WARNING_STATUS AS 预警状态, COUNT(m.ID) AS 发送消息数量 FROM BAYONET_DYNAMIC.DATAFUSION d JOIN BAYONET_BASICS.SYS_BAYONET b ON d.BAYONET_ID = b.ID LEFT JOIN BAYONET_BUSSINESS.BOAT_WARNING w ON d.CODE = w.EVENTCODE LEFT JOIN BAYONET_BUSSINESS.ZD_WARNING_TYPE wt ON w.WARNING_TYPE = wt.ID LEFT JOIN BAYONET_BUSSINESS.MESSAGE_SEND_RECORD m ON w.WARNCODE = m.WARNING_CODE WHERE d.BOAT_NAME = '苏无锡货08168' GROUP BY d.PASSTIME, b.NAME, d.BOAT_NAME, wt.NAME, w.WARNING_TIME, w.WARNING_STATUS ORDER BY d.PASSTIME DESC) WHERE ROWNUM = 1"
    }
]

# 配置Few-Shot提示模板
few_shot_example_prompt = PromptTemplate.from_template("输入问题: {input}\n生成SQL: {query}")
few_shot_prompt = FewShotPromptTemplate(
    examples=examples,
    example_prompt=few_shot_example_prompt,
    prefix="你是一个Oracle SQL专家。请根据以下表结构生成对应问题的Oracle SQL语句，且尽量符合标准语法。\n\n以下是表结构（table_info）：\n{table_info}\n\n这是一些问题及其对应的SQL样例:",
    suffix="\n---\n问题: {input}\n对应SQL: \n\n请注意：只输出SQL代码，不要包含任何解释、注释或Markdown格式",
    input_variables=["input", "table_info"],
)

# 配置分析提示模板
TEXT_PROMPT = PromptTemplate.from_template(
    """根据问题{test_question}，SQL查询结果{sql_result}，以及生成的图表类型{chart_type}，  
    请提供以下内容：  
    1. 对数据的简明分析  
    2. 从数据中得出的关键见解  
    3. 为什么选择了这种图表类型来展示数据  
    4. 根据数据可能的建议或行动计划  

    请使用简洁专业的语言回答。"""
)

# 创建SQL查询链
generate_query = create_sql_query_chain(llm=llm, db=db, prompt=SQL_PROMPT, k=5)


@app.route('/api/query', methods=['POST'])
def process_query():
    try:
        # 获取请求数据
        data = request.json
        question = data.get('question')

        if not question:
            return jsonify({'error': '请提供问题参数'}), 400

            # 生成SQL
        few_shot_sql = few_shot_prompt.format(input=question, table_info=db.table_info)
        llm_response = llm.invoke(few_shot_sql)

        # 提取SQL
        if "```sql" in llm_response:
            extracted_sql = extract_sql_from_markdown(llm_response)
        else:
            extracted_sql = llm_response.replace(";", "")

        print(f"生成的SQL: {extracted_sql}")

        # 执行SQL查询
        raw_result = db.run(extracted_sql)

        # 获取列名（从SQL或在这里定义）
        # 提取SQL中的SELECT子句中的列名，或者手动定义
        column_names = []

        # 尝试从SQL中提取列名
        import re
        match = re.search(r'SELECT\s+(.*?)\s+FROM', extracted_sql, re.IGNORECASE | re.DOTALL)
        if match:
            select_clause = match.group(1)
            # 尝试解析列及其别名
            cols = []
            for col_expr in select_clause.split(','):
                col_expr = col_expr.strip()
                # 查找AS或空格后的别名
                alias_match = re.search(r'\s+AS\s+(\w+)|(\w+)$', col_expr, re.IGNORECASE)
                if alias_match:
                    # 使用AS之后的别名，或者最后一个单词作为列名
                    cols.append(alias_match.group(1) or alias_match.group(2))
                else:
                    # 无法解析，使用整个表达式
                    cols.append(col_expr)
            column_names = cols
        else:
            # 无法从SQL提取，使用默认名称
            column_names = [f"列{i + 1}" for i in range(len(raw_result[0]))]

            # 将元组转换为字典
        converted_result = []
        for row in raw_result:
            row_dict = {}
            for i, value in enumerate(row):
                # 确保有足够的列名
                if i < len(column_names):
                    col_name = column_names[i]
                else:
                    col_name = f"列{i + 1}"
                row_dict[col_name] = value
            converted_result.append(row_dict)

        # # *** 关键步骤：转换SQL结果 ***
        serializable_result = converted_result

        print("serializable_result:", serializable_result)
        # 生成可视化前转换数据
        print("开始生成可视化...")
        try:
            output_file = visualize_sql_result(
                sql_result=raw_result,  # 使用已序列化的结果
                query=extracted_sql,
                question=question
            )
            print(f"可视化文件: {output_file}")
        except Exception as viz_error:
            print(f"可视化生成失败: {viz_error}")
            traceback.print_exc()
            output_file = None

            # 分析结果
        if output_file:
            try:
                # 从可视化推荐中获取图表类型
                df = convert_sql_result_to_dataframe(serializable_result)  # 使用已序列化的结果
                recommendation = analyze_data_for_visualization(df, extracted_sql, question)
                # 确保推荐也是可序列化的
                recommendation = serialize_row(recommendation)
                chart_type = recommendation.get("chart_type") if recommendation else "未知"

                # 生成分析
                formatted_prompt = TEXT_PROMPT.format(
                    test_question=question,
                    sql_result=json.dumps(serializable_result, ensure_ascii=False),  # 转换为JSON字符串
                    chart_type=chart_type
                )
                analysis_result = llm.invoke(formatted_prompt)
            except Exception as analysis_err:
                print(f"分析生成失败: {analysis_err}")
                traceback.print_exc()
                analysis_result = f"无法生成分析：{str(analysis_err)}"
                chart_type = None
                recommendation = None
        else:
            analysis_result = "无法生成可视化和分析"
            chart_type = None
            recommendation = None

            # 准备返回结果
        result = {
            'question': question,
            'sql': extracted_sql,
            'result': serializable_result,  # 已经序列化
            'visualization': output_file,
            'analysis': analysis_result,
            'chart_type': chart_type,
            'recommendation': recommendation
        }

        # 在返回前检查是否所有数据都可序列化
        try:
            # 尝试用json模块序列化，确保结果可以被JSON化
            json.dumps(result, cls=CustomJSONEncoder)
        except TypeError as json_err:
            print(f"警告: 结果不可JSON序列化: {json_err}")
            # 尝试修复结果，转换所有不可序列化的部分为字符串
            result = serialize_row(result)

        return jsonify(result)

    except Exception as e:
        error_trace = traceback.format_exc()
        print(f"处理查询时出错: {e}\n{error_trace}")
        return jsonify({'error': str(e), 'trace': error_trace}), 500


@app.route('/api/health', methods=['GET'])
def health_check():
    return jsonify({'status': 'healthy', 'message': 'API服务正常运行中'})


@app.route('/visualizations/<path:filename>', methods=['GET'])
def serve_visualization(filename):
    # 设置可视化文件的目录（如果它们存储在特定位置）
    viz_dir = os.path.join(os.getcwd(), 'visualizations')
    return send_from_directory(viz_dir, filename)


# 主程序入口
if __name__ == '__main__':
    # 确保只在主程序中运行，避免导入时启动
    app.run(debug=True, host='0.0.0.0', port=5005, use_reloader=False)